package com.yf.core.partitioning.impl;

import com.yf.core.partitioning.Partitioning;
import com.yf.core.partitioning.schedule_policy.OfferPolicy;
import com.yf.core.partitioning.schedule_policy.PollPolicy;
import com.yf.core.partitioning.schedule_policy.RemovePolicy;
import com.yf.core.partitioning.schedule_policy.impl.offer_policy.RoundRobinOffer;
import com.yf.core.partitioning.schedule_policy.impl.poll_policy.RoundRobinPoll;
import com.yf.core.partitioning.schedule_policy.impl.remove_policy.RoundRobinRemove;

import java.lang.reflect.InvocationTargetException;
import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author yyf
 * @date 2025/8/8 17:20
 * @description
 * 分区队列：多队列、细粒度、高性能
 * 可选任务入队策略：轮询、随机、hash、填谷
 * 可选任务出队策略：轮询、随机、削峰、线程绑定
 * 可选移除任务策略：轮询、随机、削峰
 *
 * 《分区流-动态分区队列》
 * 这是一款综合的分区队列：支持出入队调度策略执行后轮询，因而类名结尾为"Flow"，当然也可以自行设置调度策略是否支持轮询，
 * 其中线程绑定出队策略是默认不支持的，其他都是默认支持轮询
 */
public class PartiStill<T> extends AbstractCollection<T> implements BlockingQueue<T>, Partitioning<T> {
    private BlockingQueue<T>[] partitions;
    private OfferPolicy offerPolicy = new RoundRobinOffer();
    private PollPolicy pollPolicy = new RoundRobinPoll();
    private RemovePolicy removePolicy = new RoundRobinRemove();
    private Integer DEFAULT_PARTITION_NUM = 4;
    private volatile Integer capacity;
    private static final Integer DEFAULT_WAIT_TIME = 100;

    public PartiStill(Integer partitionNum, Integer capacity, Class<?> qClass, OfferPolicy offerPolicy, PollPolicy pollPolicy, RemovePolicy removePolicy) {
        this(partitionNum, capacity,qClass);
        this.offerPolicy = offerPolicy;
        this.pollPolicy = pollPolicy;
        this.removePolicy = removePolicy;
    }

    public PartiStill(Integer partitionNum, Integer capacity, Class<?> qClass) {
        partitions = new BlockingQueue[partitionNum];
        this.capacity = capacity;
        if (capacity != null) {//不为null，轮询分配容量
            final int baseCapacity = capacity / partitionNum;
            final int remainder = capacity % partitionNum;
            for (int i = 0; i < partitionNum; i++) {
                int partitionCapacity = baseCapacity + (i < remainder ? 1 : 0);
                try {
                    partitions[i] = (BlockingQueue<T>) qClass.getConstructor(int.class).newInstance(partitionCapacity);
                } catch (NoSuchMethodException e) {
                    throw new RuntimeException(e);
                } catch (InvocationTargetException e) {
                    throw new RuntimeException(e);
                } catch (InstantiationException e) {
                    throw new RuntimeException(e);
                } catch (IllegalAccessException e) {
                    throw new RuntimeException(e);
                }
            }
        } else {//为null，都为无界队列
            for (int i = 0; i < partitionNum; i++) {
                try {
                    partitions[i] = (BlockingQueue<T>) qClass.getConstructor().newInstance();
                } catch (InstantiationException e) {
                    throw new RuntimeException(e);
                } catch (IllegalAccessException e) {
                    throw new RuntimeException(e);
                } catch (InvocationTargetException e) {
                    throw new RuntimeException(e);
                } catch (NoSuchMethodException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }
    public PartiStill() {
        partitions = new BlockingQueue[DEFAULT_PARTITION_NUM];
        for (int i = 0; i < DEFAULT_PARTITION_NUM; i++) {
            partitions[i] = new LinkedBlockingQueue<>();
        }
    }


    @Override
    public BlockingQueue<T>[] getPartitions() {
        return partitions;
    }

    @Override
    public OfferPolicy getOfferPolicy() {
        return null;
    }

    @Override
    public PollPolicy getPollPolicy() {
        return null;
    }

    @Override
    public RemovePolicy getRemovePolicy() {
        return null;
    }

    @Override
    public boolean add(T t) {
        return partitions[offerPolicy.selectPartition(partitions, t)].add(t);
    }

    @Override
    public boolean offer(T t) {
        return partitions[offerPolicy.selectPartition(partitions, t)].offer(t);
    }

    @Override
    public T remove() {
        return partitions[removePolicy.selectPartition(partitions)].remove();
    }

    @Override
    public T poll() {
        return partitions[pollPolicy.selectPartition(partitions)].poll();
    }

    @Override
    public T element() {
        return partitions[pollPolicy.selectPartition(partitions)].peek();
    }

    @Override
    public T peek() {
        return partitions[pollPolicy.selectPartition(partitions)].peek();
    }

    @Override
    public void put(T t) throws InterruptedException {
        partitions[offerPolicy.selectPartition(partitions, t)].put(t);
    }

    @Override
    public boolean offer(T t, long timeout, TimeUnit unit) throws InterruptedException {
        return partitions[pollPolicy.selectPartition(partitions)].offer(t,timeout, unit);
    }

    @Override
    public T take() throws InterruptedException {
        return partitions[pollPolicy.selectPartition(partitions)].take();
    }

    @Override
    public T poll(long timeout, TimeUnit unit) throws InterruptedException {
        return partitions[pollPolicy.selectPartition(partitions)].poll(timeout, unit);
    }

    @Override
    public int remainingCapacity() {
        int capacity = 0;
        for(BlockingQueue<T> q: partitions){
            capacity += q.remainingCapacity();
        }
        return capacity;
    }



    @Override
    public int size() {
        int totalSize = 0;
        for (BlockingQueue<T> partition : partitions) {
            totalSize += partition.size();
        }
        return totalSize;
    }




    @Override
    public Iterator<T> iterator() {
        List<T> allElements = new ArrayList<>();
        for (BlockingQueue<T> partition : partitions) {
            allElements.addAll(partition);
        }
        return allElements.iterator();
    }


    @Override
    public int drainTo(Collection<? super T> c) {
        int count = 0;
        for (BlockingQueue<T> partition : partitions) {
            count += partition.drainTo(c);
        }
        return count;
    }

    @Override
    public int drainTo(Collection<? super T> c, int maxElements) {
        int count = 0;
        for (BlockingQueue<T> partition : partitions) {
            if (count >= maxElements) {
                break;
            }
            count += partition.drainTo(c, maxElements - count);
        }
        return count;
    }
}
